package cipher

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"errors"
	"io"
	"os"
)

var ivDefValue = []byte("01020304050607080102030405060708")


/**
将key长度变为16字节
 */
func checkKey(key []byte,expectLen int) ([]byte ,error) {
	KEYLEN:=8
	switch expectLen {
	case 8,16,24:
		KEYLEN=expectLen
	}
	if key==nil{
		return nil,errors.New("输入key==nil")
	}
	keyLen:=len(key)
	if keyLen>=KEYLEN{
		return key[:KEYLEN],nil
	} else {
		repeatCount:=KEYLEN/keyLen+1
		rKey:=bytes.Repeat(key,repeatCount)
		return rKey[:KEYLEN],nil
	}
}

/**
Key长度随便输，会自动增减成16字节
加密后的长度为原长度的len/16+16字节
 */
func AesEncrypt(plaintext []byte, key []byte) ([]byte, error) {
	newKey,err:=checkKey(key,16)
	if err != nil {
		return nil, errors.New("invalid key")
	}
	block, err := aes.NewCipher(newKey)
	if err != nil {
		return nil, errors.New("invalid encrypt key")
	}
	blockSize := block.BlockSize()
	plaintext = PKCS5Padding(plaintext, blockSize)
	iv := ivDefValue[:blockSize]
	blockMode := cipher.NewCBCEncrypter(block, iv)

	ciphertext := make([]byte, len(plaintext))
	blockMode.CryptBlocks(ciphertext, plaintext)

	return ciphertext, nil
}

func AesDecrypt(ciphertext []byte, key []byte) ([]byte, error) {

	newKey,err:=checkKey(key,16)
	if err != nil {
		return nil, errors.New("checkKey:invalid key")
	}
	block, err := aes.NewCipher(newKey)
	if err != nil {
		return nil, errors.New("invalid decrypt key")
	}

	blockSize := block.BlockSize()

	if len(ciphertext) < blockSize {
		return nil, errors.New("ciphertext too short")
	}

	iv := ivDefValue[:blockSize]
	if len(ciphertext)%blockSize != 0 {
		return nil, errors.New("ciphertext is not a multiple of the block size")
	}

	blockModel := cipher.NewCBCDecrypter(block, iv)

	plaintext := make([]byte, len(ciphertext))
	blockModel.CryptBlocks(plaintext, ciphertext)
	plaintext = PKCS5UnPadding(plaintext)

	return plaintext, nil
}


func PKCS5Padding(src []byte, blockSize int) []byte {
	padding := blockSize - len(src)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(src, padtext...)
}

func PKCS5UnPadding(src []byte) []byte {
	length := len(src)
	unpadding := int(src[length-1])
	return src[:(length - unpadding)]
}


/**
AES加密文件用CTR，不改变长度，原来多长，处理完成后就多长
加解密同一个函数
 */
func AesCtr(src, key []byte) ([]byte, error) {
	newKey,err:=checkKey(key,16)
	if err != nil {
		return nil, errors.New("checkKey:invalid key")
	}
	//1. 创建cipher.Block接口
	block, err := aes.NewCipher(newKey)
	if err != nil {
		return nil, err
	}
	//2. 创建分组模式，在crypto/cipher包中
	stream := cipher.NewCTR(block, ivDefValue[:block.BlockSize()])
	//3. 异或
	dst := make([]byte, len(src))
	stream.XORKeyStream(dst, src)

	return dst, nil
}


/**
AES加密文件用CTR，不改变文件长度
加解密同一个函数
 */
func AesCtrFile(srcFile,destFile string,key []byte) error {
	newKey,err:=checkKey(key,16)
	if err != nil {
		return  errors.New("checkKey:invalid key")
	}
	readFile, err := os.Open(srcFile)
	outFile, err := os.OpenFile(destFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		panic(err)
	}
	defer readFile.Close()
	defer outFile.Close()
	block, err := aes.NewCipher(newKey)
	blockSize:=block.BlockSize()
	if err != nil {
		panic(err)
	}
	writer := &cipher.StreamWriter{S: cipher.NewCTR(block, ivDefValue[:blockSize]), W: outFile}
	if _, err := io.Copy(writer, readFile); err != nil {
		return err
	}
	return nil
}


